/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "LinkedListOIT.h"
#include "SceneBase.h"
#include "DXHelper.h"

LinkedListOIT g_linkedListOIT;

/**
*/
LinkedListOIT::LinkedListOIT() : m_pFragmentLink( NULL ), m_pFragmentLinkSRV( NULL ), m_pFragmentLinkUAV( NULL ),
                                 m_pStartOffsetBuffer( NULL ), m_pStartOffsetSRV( NULL ), m_pStartOffsetUAV( NULL ),
                                 m_pDebugPS( NULL ), m_pDepthStencilState( NULL ), m_pDebugRS( NULL ), m_pRS( NULL ), m_pAlphaBlendBS( NULL ), 
                                 m_pPS_CB( NULL ), m_blendMode( BLEND_OIT )
{
}

/**
*/
LinkedListOIT::~LinkedListOIT()
{
}

/**
*/
HRESULT LinkedListOIT::OnCreateDevice( ID3D11Device* pDevice )
{
    HRESULT hr;

    // Create depth/stencil state
    // turn off z write
    D3D11_DEPTH_STENCIL_DESC DSDesc;
    ZeroMemory( &DSDesc, sizeof( D3D11_DEPTH_STENCIL_DESC ) );
    DSDesc.DepthEnable = FALSE;
    DSDesc.DepthFunc = D3D11_COMPARISON_LESS;
    DSDesc.StencilEnable = FALSE;
    DSDesc.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ZERO;
    V_RETURN( pDevice->CreateDepthStencilState( &DSDesc, &m_pDepthStencilState ) );

    // Constant Buffer
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;
    Desc.ByteWidth = sizeof( PS_CB );
    V_RETURN( pDevice->CreateBuffer( &Desc, NULL, &m_pPS_CB ) );

    V_RETURN( m_storeFragments.OnCreateDevice( pDevice ) );
    V_RETURN( m_sortFragmentsAndRender.OnCreateDevice( pDevice ) );

    // Debug pixel shader
    ID3DBlob* pBlob = NULL;
    V_RETURN( CompileShaderFromFile( L"Scene.hlsl", "ScenePS", "ps_5_0", &pBlob ) );
    V_RETURN( pDevice->CreatePixelShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pDebugPS ) );
    SAFE_RELEASE( pBlob );

    // Debug rasterizer state
    {
        D3D11_RASTERIZER_DESC drd = {
            D3D11_FILL_WIREFRAME, //D3D11_FILL_MODE FillMode;
            D3D11_CULL_BACK,//D3D11_CULL_MODE CullMode;
            FALSE, //BOOL FrontCounterClockwise;
            0, //INT DepthBias;
            0.0f,//FLOAT DepthBiasClamp;
            0.0f,//FLOAT SlopeScaledDepthBias;
            TRUE,//BOOL DepthClipEnable;
            FALSE,//BOOL ScissorEnable;
            TRUE,//BOOL MultisampleEnable;
            FALSE//BOOL AntialiasedLineEnable;        
        };
        pDevice->CreateRasterizerState( &drd, &m_pDebugRS );
    }

    // rasterizer state      
    {
        D3D11_RASTERIZER_DESC drd = {
            D3D11_FILL_SOLID, //D3D11_FILL_MODE FillMode;
            D3D11_CULL_NONE,//D3D11_CULL_MODE CullMode;
            FALSE, //BOOL FrontCounterClockwise;
            0, //INT DepthBias;
            0.0f,//FLOAT DepthBiasClamp;
            0.0f,//FLOAT SlopeScaledDepthBias;
            TRUE,//BOOL DepthClipEnable;
            FALSE,//BOOL ScissorEnable;
            TRUE,//BOOL MultisampleEnable;
            FALSE//BOOL AntialiasedLineEnable;        
        };
        pDevice->CreateRasterizerState( &drd, &m_pRS );
    }

    // alpha blending enabled blend state
    D3D11_BLEND_DESC BSDesc;
    ZeroMemory( &BSDesc, sizeof( D3D11_BLEND_DESC ) );

    BSDesc.RenderTarget[0].BlendEnable = TRUE;
    BSDesc.RenderTarget[0].SrcBlend = D3D11_BLEND_SRC_ALPHA;
    BSDesc.RenderTarget[0].DestBlend = D3D11_BLEND_INV_SRC_ALPHA;
    BSDesc.RenderTarget[0].BlendOp = D3D11_BLEND_OP_ADD;
    BSDesc.RenderTarget[0].SrcBlendAlpha = D3D11_BLEND_ONE;
    BSDesc.RenderTarget[0].DestBlendAlpha = D3D11_BLEND_ZERO;
    BSDesc.RenderTarget[0].BlendOpAlpha = D3D11_BLEND_OP_ADD;
    BSDesc.RenderTarget[0].RenderTargetWriteMask = 0x0F;

    V_RETURN( pDevice->CreateBlendState( &BSDesc, &m_pAlphaBlendBS ) );

    return S_OK;
}

/**
*/
HRESULT LinkedListOIT::OnResizedSwapChain( const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc, ID3D11Device* pDevice )
{
    HRESULT hr;

    // Create Fragment and Link buffer.
    V_RETURN( CreateFragmentLink( pBackBufferSurfaceDesc, pDevice ) );

    // Create Start Offset buffer
    V_RETURN( CreateStartOffset( pBackBufferSurfaceDesc, pDevice ) );

    return hr;
}

/**
*/
HRESULT LinkedListOIT::CreateFragmentLink( const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc, ID3D11Device* pDevice )
{
    HRESULT hr;

    const int NUM_ELEMENTS = pBackBufferSurfaceDesc->Width * pBackBufferSurfaceDesc->Height * 8;

    // create buffer
    D3D11_BUFFER_DESC descBuf;
    memset( &descBuf, 0, sizeof( descBuf ) );
    descBuf.StructureByteStride = sizeof( FragmentLink );
    descBuf.ByteWidth = NUM_ELEMENTS * descBuf.StructureByteStride;
    descBuf.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_STRUCTURED;
    descBuf.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
    V_RETURN( pDevice->CreateBuffer( &descBuf, NULL, &m_pFragmentLink ) );

    // create UAV
    D3D11_UNORDERED_ACCESS_VIEW_DESC descUAV;
    memset( &descUAV, 0, sizeof( descUAV ) );
    descUAV.ViewDimension = D3D11_UAV_DIMENSION_BUFFER;
    descUAV.Buffer.FirstElement = 0;
    descUAV.Format = DXGI_FORMAT_UNKNOWN;
    descUAV.Buffer.NumElements = NUM_ELEMENTS;
    descUAV.Buffer.Flags = D3D11_BUFFER_UAV_FLAG_COUNTER;
    V_RETURN( pDevice->CreateUnorderedAccessView( m_pFragmentLink, &descUAV, &m_pFragmentLinkUAV ) );

    // create SRV
    D3D11_SHADER_RESOURCE_VIEW_DESC descSRV;
    descSRV.ViewDimension = D3D11_SRV_DIMENSION_BUFFER;
    descSRV.Buffer.FirstElement = 0;
    descSRV.Format = DXGI_FORMAT_UNKNOWN;
    descSRV.Buffer.NumElements = NUM_ELEMENTS;
    V_RETURN( pDevice->CreateShaderResourceView( m_pFragmentLink, &descSRV, &m_pFragmentLinkSRV ) );

    return hr;
}

/**
*/
HRESULT LinkedListOIT::CreateStartOffset( const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc, ID3D11Device* pDevice )
{
    HRESULT hr;

    const int NUM_ELEMENTS = pBackBufferSurfaceDesc->Width * pBackBufferSurfaceDesc->Height;

    // create buffer
    D3D11_BUFFER_DESC descBuf;
    memset( &descBuf, 0, sizeof( descBuf ) );
    descBuf.StructureByteStride = sizeof( unsigned int );
    descBuf.ByteWidth = NUM_ELEMENTS * descBuf.StructureByteStride;
    descBuf.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_ALLOW_RAW_VIEWS;
    descBuf.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
    V_RETURN( pDevice->CreateBuffer( &descBuf, NULL, &m_pStartOffsetBuffer ) );

    // create UAV
    D3D11_UNORDERED_ACCESS_VIEW_DESC descUAV;
    memset( &descUAV, 0, sizeof( descUAV ) );
    descUAV.ViewDimension = D3D11_UAV_DIMENSION_BUFFER;
    descUAV.Buffer.FirstElement = 0;
    descUAV.Format = DXGI_FORMAT_R32_TYPELESS;
    descUAV.Buffer.NumElements = NUM_ELEMENTS;
    descUAV.Buffer.Flags = D3D11_BUFFER_UAV_FLAG_RAW;
    V_RETURN( pDevice->CreateUnorderedAccessView( m_pStartOffsetBuffer, &descUAV, &m_pStartOffsetUAV ) );

    // create SRV
    D3D11_SHADER_RESOURCE_VIEW_DESC descSRV;
    descSRV.ViewDimension = D3D11_SRV_DIMENSION_BUFFER;
    descSRV.Buffer.FirstElement = 0;
    descSRV.Format = DXGI_FORMAT_R32_UINT;
    descSRV.Buffer.NumElements = NUM_ELEMENTS;
    V_RETURN( pDevice->CreateShaderResourceView( m_pStartOffsetBuffer, &descSRV, &m_pStartOffsetSRV ) );

    return hr;
}

/**
*/
void LinkedListOIT::OnReleasingSwapChain()
{
    // Fragment And Link
    SAFE_RELEASE( m_pFragmentLink );
    SAFE_RELEASE( m_pFragmentLinkSRV );
    SAFE_RELEASE( m_pFragmentLinkUAV );

    // Start Offset
    SAFE_RELEASE( m_pStartOffsetBuffer );
    SAFE_RELEASE( m_pStartOffsetSRV );
    SAFE_RELEASE( m_pStartOffsetUAV );    
}

/**
*/
void LinkedListOIT::OnDestroyDevice()
{
    SAFE_RELEASE( m_pDepthStencilState );
    SAFE_RELEASE( m_pDebugRS );
    SAFE_RELEASE( m_pRS );
    SAFE_RELEASE( m_pPS_CB );
    SAFE_RELEASE( m_pDebugPS );
    SAFE_RELEASE( m_pAlphaBlendBS );

    m_storeFragments.OnDestroyDevice();
    m_sortFragmentsAndRender.OnDestroyDevice();
}

/**
*/
void LinkedListOIT::OnRender( ID3D11DeviceContext* pD3DContext, ID3D11Device* pDevice, const DXGI_SURFACE_DESC* pBackBufferDesc,
                              SceneBase* pScene, D3DXMATRIXA16* pmViewProj, const D3DXVECTOR3* pvCamEyePos, ID3D11RenderTargetView* pRTV, ID3D11DepthStencilView* pDSV )
{
    HRESULT hr;

    // Save rasterizer state
    ID3D11RasterizerState* pOldState;
    pD3DContext->RSGetState( &pOldState );
    pD3DContext->RSSetState( m_pRS );
    
    // Save depth/stencil state
    ID3D11DepthStencilState* pDepthStencilStateStored = NULL;
    UINT stencilRef;
    pD3DContext->OMGetDepthStencilState( &pDepthStencilStateStored, &stencilRef );

    if( m_blendMode == BLEND_OIT )
    {
        // Setup the constant buffer
        D3D11_MAPPED_SUBRESOURCE MappedResource;
        V( pD3DContext->Map( m_pPS_CB, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
        PS_CB* pPS_CB = ( PS_CB* )MappedResource.pData;  
        pPS_CB->nFrameWidth = pBackBufferDesc->Width;
        pPS_CB->nFrameHeight = pBackBufferDesc->Height;
        pD3DContext->Unmap( m_pPS_CB, 0 );
        pD3DContext->PSSetConstantBuffers( 0, 1, &m_pPS_CB );

        // Store the fragment linked-list into the structured buffer.
        m_storeFragments.OnRender( pD3DContext, pDevice, pBackBufferDesc, pScene, pmViewProj, pvCamEyePos, pRTV, pDSV, 
            m_pFragmentLinkUAV, m_pStartOffsetUAV, m_pDepthStencilState, m_pFragmentLink, m_pStartOffsetBuffer );
            
        // Sort and render linked-list fragments into the frame buffer.
        pD3DContext->OMSetRenderTargets( 1, &pRTV, pDSV );
        float ClearColor[4] = { 1.0f, 0.0f, 1.0f, 1.0f }; 

        pD3DContext->OMSetDepthStencilState( m_pDepthStencilState, 0 );
        pD3DContext->ClearRenderTargetView( pRTV, ClearColor );
        
        m_sortFragmentsAndRender.OnRender( pD3DContext, pDevice, pBackBufferDesc, m_pFragmentLinkSRV, m_pStartOffsetSRV );
    }
    else if( m_blendMode == BLEND_REGULAR )
    {
        pD3DContext->OMSetRenderTargets( 1, &pRTV, pDSV );

        ID3D11BlendState* pBlendStateStored11;
        float fBlendFactorStored11[4];
        UINT nSampleMaskStored11;

        pD3DContext->OMGetBlendState( &pBlendStateStored11, fBlendFactorStored11, &nSampleMaskStored11 );

        // enable alphablending
        float BlendFactor[4] = { 0, 0, 0, 0 };
        pD3DContext->OMSetBlendState( m_pAlphaBlendBS, BlendFactor, 0xFFFFFFFF );

        // disable depth stencil tests
        pD3DContext->OMSetDepthStencilState( m_pDepthStencilState, 0 );

        // Kind of ugly. Since pixel shader is not set inside pScene->OnRender, set it here.
        pD3DContext->PSSetShader( m_pDebugPS, NULL, 0 );

        pScene->OnRender( pD3DContext, pmViewProj, pvCamEyePos );

        pD3DContext->OMSetBlendState( pBlendStateStored11, fBlendFactorStored11, nSampleMaskStored11 );

        SAFE_RELEASE( pBlendStateStored11 );
    }
    else 
    {
        assert( 0 );
    }

    // Restore saved render state.
    pD3DContext->RSSetState( pOldState );
    SAFE_RELEASE( pOldState );

    // Restore saved depth/stencil state.
    pD3DContext->OMSetDepthStencilState( pDepthStencilStateStored, stencilRef );
    SAFE_RELEASE( pDepthStencilStateStored );
}
